{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2KROuTZVuhrp"
   },
   "source": [
    "##### Copyright 2020 The TensorFlow Authors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "id": "6aEVQQ403kzs"
   },
   "outputs": [],
   "source": [
    "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "P5VpOpSivqgv"
   },
   "source": [
    "# Customizing what happens in `fit()`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vApNeEfvLLc4"
   },
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/keras-team/keras-io/blob/master/tf/customizing_what_happens_in_fit.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/keras-team/keras-io/blob/master/guides/customizing_what_happens_in_fit.py\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://storage.googleapis.com/tensorflow_docs/keras-io/tf/customizing_what_happens_in_fit.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TihRiHIKeeIz"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "When you're doing supervised learning, you can use `fit()` and everything works\n",
    "smoothly.\n",
    "\n",
    "When you need to write your own training loop from scratch, you can use the\n",
    "`GradientTape` and take control of every little detail.\n",
    "\n",
    "But what if you need a custom training algorithm, but you still want to benefit from\n",
    "the convenient features of `fit()`, such as callbacks, built-in distribution support,\n",
    "or step fusing?\n",
    "\n",
    "A core principle of Keras is **progressive disclosure of complexity**. You should\n",
    "always be able to get into lower-level workflows in a gradual way. You shouldn't fall\n",
    "off a cliff if the high-level functionality doesn't exactly match your use case. You\n",
    "should be able to gain more control over the small details while retaing a\n",
    "commensurate amount of high-level convenience.\n",
    "\n",
    "When you need to customize what `fit()` does, you should **override the training step\n",
    "function of the `Model` class**. This is the function that is called by `fit()` for\n",
    "every batch of data. You will then be able to call `fit()` as usual -- and it will be\n",
    "running your own learning algorithm.\n",
    "\n",
    "Note that this pattern does not prevent you from building models with the Functional\n",
    "API. You can do this whether you're building `Sequential` models, Functional API\n",
    "models, or subclassed models.\n",
    "\n",
    "Let's see how that works."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XFRryV6yxq2Z"
   },
   "source": [
    "## Setup\n",
    "Requires TensorFlow 2.2 or later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code",
    "id": "BxGJZEXaWrLM"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1yZO4J3zyOfz"
   },
   "source": [
    "## A first simple example\n",
    "\n",
    "Let's start from a simple example:\n",
    "\n",
    "- We create a new class that subclasses `keras.Model`.\n",
    "- We just override the method `train_step(self, data)`.\n",
    "- We return a dictionary mapping metric names (including the loss) to their current\n",
    "value.\n",
    "\n",
    "The input argument `data` is what gets passed to fit as training data:\n",
    "\n",
    "- If you pass Numpy arrays, by calling `fit(x, y, ...)`, then `data` will be the tuple\n",
    "`(x, y)`\n",
    "- If you pass a `tf.data.Dataset`, by calling `fit(dataset, ...)`, then `data` will be\n",
    "what gets yielded by `dataset` at each batch.\n",
    "\n",
    "In the body of the `train_step` method, we implement a regular training update,\n",
    "similar to what you are already familiar with. Importantly, **we compute the loss via\n",
    "`self.compiled_loss`**, which wraps the loss(es) function(s) that were passed to\n",
    "`compile()`.\n",
    "\n",
    "Similarly, we call `self.compiled_metrics.update_state(y, y_pred)` to update the state\n",
    "of the metrics that were passed in `compile()`, and we query results from\n",
    "`self.metrics` at the end to retrieve their current value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code",
    "id": "sg0aNp6yuNUs"
   },
   "outputs": [],
   "source": [
    "class CustomModel(keras.Model):\n",
    "    def train_step(self, data):\n",
    "        # Unpack the data. Its structure depends on your model and\n",
    "        # on what you pass to `fit()`.\n",
    "        x, y = data\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            y_pred = self(x, training=True)  # Forward pass\n",
    "            # Compute the loss value\n",
    "            # (the loss function is configured in `compile()`)\n",
    "            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)\n",
    "\n",
    "        # Compute gradients\n",
    "        trainable_vars = self.trainable_variables\n",
    "        gradients = tape.gradient(loss, trainable_vars)\n",
    "        # Update weights\n",
    "        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n",
    "        # Update metrics (includes the metric that tracks the loss)\n",
    "        self.compiled_metrics.update_state(y, y_pred)\n",
    "        # Return a dict mapping metric names to current value\n",
    "        return {m.name: m.result() for m in self.metrics}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YEdOFRbXmA4d"
   },
   "source": [
    "Let's try this out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code",
    "id": "1wDUe4ReTaVi"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Construct and compile an instance of CustomModel\n",
    "inputs = keras.Input(shape=(32,))\n",
    "outputs = keras.layers.Dense(1)(inputs)\n",
    "model = CustomModel(inputs, outputs)\n",
    "model.compile(optimizer=\"adam\", loss=\"mse\", metrics=[\"mae\"])\n",
    "\n",
    "# Just use `fit` as usual\n",
    "x = np.random.random((1000, 32))\n",
    "y = np.random.random((1000, 1))\n",
    "model.fit(x, y, epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tQSwBvcGIeZk"
   },
   "source": [
    "## Going lower-level\n",
    "\n",
    "Naturally, you could just skip passing a loss function in `compile()`, and instead do\n",
    "everything *manually* in `train_step`. Likewise for metrics. Here's a lower-level\n",
    "example, that only uses `compile()` to configure the optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code",
    "id": "9UnwB6gdESVw"
   },
   "outputs": [],
   "source": [
    "mae_metric = keras.metrics.MeanAbsoluteError(name=\"mae\")\n",
    "loss_tracker = keras.metrics.Mean(name=\"loss\")\n",
    "\n",
    "\n",
    "class CustomModel(keras.Model):\n",
    "    def train_step(self, data):\n",
    "        x, y = data\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            y_pred = self(x, training=True)  # Forward pass\n",
    "            # Compute our own loss\n",
    "            loss = keras.losses.mean_squared_error(y, y_pred)\n",
    "\n",
    "        # Compute gradients\n",
    "        trainable_vars = self.trainable_variables\n",
    "        gradients = tape.gradient(loss, trainable_vars)\n",
    "\n",
    "        # Update weights\n",
    "        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n",
    "\n",
    "        # Compute our own metrics\n",
    "        loss_tracker.update_state(loss)\n",
    "        mae_metric.update_state(y, y_pred)\n",
    "        return {\"loss\": loss_tracker.result(), \"mae\": mae_metric.result()}\n",
    "\n",
    "\n",
    "# Construct an instance of CustomModel\n",
    "inputs = keras.Input(shape=(32,))\n",
    "outputs = keras.layers.Dense(1)(inputs)\n",
    "model = CustomModel(inputs, outputs)\n",
    "\n",
    "# We don't passs a loss or metrics here.\n",
    "model.compile(optimizer=\"adam\")\n",
    "\n",
    "# Just use `fit` as usual -- you can use callbacks, etc.\n",
    "x = np.random.random((1000, 32))\n",
    "y = np.random.random((1000, 1))\n",
    "model.fit(x, y, epochs=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WN0qnQacU9u2"
   },
   "source": [
    "Note that with this setup, you will need to manually call `reset_states()` on your\n",
    "metrics after each epoch, or between training and evaluation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "fnMF4QYQFNj1"
   },
   "source": [
    "## Supporting `sample_weight` & `class_weight`\n",
    "\n",
    "You may have noticed that our first basic example didn't make any mention of sample\n",
    "weighting. If you want to support the `fit()` arguments `sample_weight` and\n",
    "`class_weight`, you'd simply do the following:\n",
    "\n",
    "- Unpack `sample_weight` from the `data` argument\n",
    "- Pass it to `compiled_loss` & `compiled_metrics` (of course, you could also just apply\n",
    "it manually if you don't rely on `compile()` for losses & metrics)\n",
    "- That's it. That's the list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code",
    "id": "tE4yrX22rlL4"
   },
   "outputs": [],
   "source": [
    "class CustomModel(keras.Model):\n",
    "    def train_step(self, data):\n",
    "        # Unpack the data. Its structure depends on your model and\n",
    "        # on what you pass to `fit()`.\n",
    "        if len(data) == 3:\n",
    "            x, y, sample_weight = data\n",
    "        else:\n",
    "            x, y = data\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            y_pred = self(x, training=True)  # Forward pass\n",
    "            # Compute the loss value.\n",
    "            # The loss function is configured in `compile()`.\n",
    "            loss = self.compiled_loss(\n",
    "                y,\n",
    "                y_pred,\n",
    "                sample_weight=sample_weight,\n",
    "                regularization_losses=self.losses,\n",
    "            )\n",
    "\n",
    "        # Compute gradients\n",
    "        trainable_vars = self.trainable_variables\n",
    "        gradients = tape.gradient(loss, trainable_vars)\n",
    "\n",
    "        # Update weights\n",
    "        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n",
    "\n",
    "        # Update the metrics.\n",
    "        # Metrics are configured in `compile()`.\n",
    "        self.compiled_metrics.update_state(y, y_pred, sample_weight=sample_weight)\n",
    "\n",
    "        # Return a dict mapping metric names to current value.\n",
    "        # Note that it will include the loss (tracked in self.metrics).\n",
    "        return {m.name: m.result() for m in self.metrics}\n",
    "\n",
    "\n",
    "# Construct and compile an instance of CustomModel\n",
    "inputs = keras.Input(shape=(32,))\n",
    "outputs = keras.layers.Dense(1)(inputs)\n",
    "model = CustomModel(inputs, outputs)\n",
    "model.compile(optimizer=\"adam\", loss=\"mse\", metrics=[\"mae\"])\n",
    "\n",
    "# You can now use sample_weight argument\n",
    "x = np.random.random((1000, 32))\n",
    "y = np.random.random((1000, 1))\n",
    "sw = np.random.random((1000, 1))\n",
    "model.fit(x, y, sample_weight=sw, epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "j0uOhTfBjhYX"
   },
   "source": [
    "## Providing your own evaluation step\n",
    "\n",
    "What if you want to do the same for calls to `model.evaluate()`? Then you would\n",
    "override `test_step` in exactly the same way. Here's what it looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code",
    "id": "vaogkBppfg2t"
   },
   "outputs": [],
   "source": [
    "class CustomModel(keras.Model):\n",
    "    def test_step(self, data):\n",
    "        # Unpack the data\n",
    "        x, y = data\n",
    "        # Compute predictions\n",
    "        y_pred = self(x, training=False)\n",
    "        # Updates the metrics tracking the loss\n",
    "        self.compiled_loss(y, y_pred, regularization_losses=self.losses)\n",
    "        # Update the metrics.\n",
    "        self.compiled_metrics.update_state(y, y_pred)\n",
    "        # Return a dict mapping metric names to current value.\n",
    "        # Note that it will include the loss (tracked in self.metrics).\n",
    "        return {m.name: m.result() for m in self.metrics}\n",
    "\n",
    "\n",
    "# Construct an instance of CustomModel\n",
    "inputs = keras.Input(shape=(32,))\n",
    "outputs = keras.layers.Dense(1)(inputs)\n",
    "model = CustomModel(inputs, outputs)\n",
    "model.compile(loss=\"mse\", metrics=[\"mae\"])\n",
    "\n",
    "# Evaluate with our custom test_step\n",
    "x = np.random.random((1000, 32))\n",
    "y = np.random.random((1000, 1))\n",
    "model.evaluate(x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xiE4ZsCtjI9B"
   },
   "source": [
    "## Wrapping up: an end-to-end GAN example\n",
    "\n",
    "Let's walk through an end-to-end example that leverages everything you just learned.\n",
    "\n",
    "Let's consider:\n",
    "\n",
    "- A generator network meant to generate 28x28x1 images.\n",
    "- A discriminator network meant to classify 28x28x1 images into two classes (\"fake\" and\n",
    "\"real\").\n",
    "- One optimizer for each.\n",
    "- A loss function to train the discriminator.\n",
    ""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code",
    "id": "jyyxuepxgMuF"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras import layers\n",
    "\n",
    "# Create the discriminator\n",
    "discriminator = keras.Sequential(\n",
    "    [\n",
    "        keras.Input(shape=(28, 28, 1)),\n",
    "        layers.Conv2D(64, (3, 3), strides=(2, 2), padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2D(128, (3, 3), strides=(2, 2), padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.GlobalMaxPooling2D(),\n",
    "        layers.Dense(1),\n",
    "    ],\n",
    "    name=\"discriminator\",\n",
    ")\n",
    "\n",
    "# Create the generator\n",
    "latent_dim = 128\n",
    "generator = keras.Sequential(\n",
    "    [\n",
    "        keras.Input(shape=(latent_dim,)),\n",
    "        # We want to generate 128 coefficients to reshape into a 7x7x128 map\n",
    "        layers.Dense(7 * 7 * 128),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Reshape((7, 7, 128)),\n",
    "        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2D(1, (7, 7), padding=\"same\", activation=\"sigmoid\"),\n",
    "    ],\n",
    "    name=\"generator\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cxFFOFm7xbCM"
   },
   "source": [
    "Here's a feature-complete GAN class, overriding `compile()` to use its own signature,\n",
    "and implementing the entire GAN algorithm in 17 lines in `train_step`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code",
    "id": "wvS2v5pvGM7h"
   },
   "outputs": [],
   "source": [
    "class GAN(keras.Model):\n",
    "    def __init__(self, discriminator, generator, latent_dim):\n",
    "        super(GAN, self).__init__()\n",
    "        self.discriminator = discriminator\n",
    "        self.generator = generator\n",
    "        self.latent_dim = latent_dim\n",
    "\n",
    "    def compile(self, d_optimizer, g_optimizer, loss_fn):\n",
    "        super(GAN, self).compile()\n",
    "        self.d_optimizer = d_optimizer\n",
    "        self.g_optimizer = g_optimizer\n",
    "        self.loss_fn = loss_fn\n",
    "\n",
    "    def train_step(self, real_images):\n",
    "        if isinstance(real_images, tuple):\n",
    "            real_images = real_images[0]\n",
    "        # Sample random points in the latent space\n",
    "        batch_size = tf.shape(real_images)[0]\n",
    "        random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))\n",
    "\n",
    "        # Decode them to fake images\n",
    "        generated_images = self.generator(random_latent_vectors)\n",
    "\n",
    "        # Combine them with real images\n",
    "        combined_images = tf.concat([generated_images, real_images], axis=0)\n",
    "\n",
    "        # Assemble labels discriminating real from fake images\n",
    "        labels = tf.concat(\n",
    "            [tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0\n",
    "        )\n",
    "        # Add random noise to the labels - important trick!\n",
    "        labels += 0.05 * tf.random.uniform(tf.shape(labels))\n",
    "\n",
    "        # Train the discriminator\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = self.discriminator(combined_images)\n",
    "            d_loss = self.loss_fn(labels, predictions)\n",
    "        grads = tape.gradient(d_loss, self.discriminator.trainable_weights)\n",
    "        self.d_optimizer.apply_gradients(\n",
    "            zip(grads, self.discriminator.trainable_weights)\n",
    "        )\n",
    "\n",
    "        # Sample random points in the latent space\n",
    "        random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))\n",
    "\n",
    "        # Assemble labels that say \"all real images\"\n",
    "        misleading_labels = tf.zeros((batch_size, 1))\n",
    "\n",
    "        # Train the generator (note that we should *not* update the weights\n",
    "        # of the discriminator)!\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = self.discriminator(self.generator(random_latent_vectors))\n",
    "            g_loss = self.loss_fn(misleading_labels, predictions)\n",
    "        grads = tape.gradient(g_loss, self.generator.trainable_weights)\n",
    "        self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))\n",
    "        return {\"d_loss\": d_loss, \"g_loss\": g_loss}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FGTUQysnjlsX"
   },
   "source": [
    "Let's test-drive it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code",
    "id": "mPJp4mErKaq1"
   },
   "outputs": [],
   "source": [
    "# Prepare the dataset. We use both the training & test MNIST digits.\n",
    "batch_size = 64\n",
    "(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()\n",
    "all_digits = np.concatenate([x_train, x_test])\n",
    "all_digits = all_digits.astype(\"float32\") / 255.0\n",
    "all_digits = np.reshape(all_digits, (-1, 28, 28, 1))\n",
    "dataset = tf.data.Dataset.from_tensor_slices(all_digits)\n",
    "dataset = dataset.shuffle(buffer_size=1024).batch(batch_size)\n",
    "\n",
    "gan = GAN(discriminator=discriminator, generator=generator, latent_dim=latent_dim)\n",
    "gan.compile(\n",
    "    d_optimizer=keras.optimizers.Adam(learning_rate=0.0003),\n",
    "    g_optimizer=keras.optimizers.Adam(learning_rate=0.0003),\n",
    "    loss_fn=keras.losses.BinaryCrossentropy(from_logits=True),\n",
    ")\n",
    "\n",
    "# To limit execution time, we only train on 100 batches. You can train on\n",
    "# the entire dataset. You will need about 20 epochs to get nice results.\n",
    "gan.fit(dataset.take(100), epochs=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZTGksVTeL0qk"
   },
   "source": [
    "The idea behind deep learning are simple, so why should their implementation be painful?"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "customizing_what_happens_in_fit.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}